package com.cqs.example.websocket.h5;

import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketMessage;
import org.springframework.web.socket.WebSocketSession;

import java.io.IOException;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author lixiaowen
 * @date 2020/07/09
 */
@Component
@Slf4j
public class WebSocketHandler implements org.springframework.web.socket.WebSocketHandler {

    /**
     * key: userName
     * value: WebSocketSession
     */
    private static final Map<String, Set<WebSocketSession>> USERS = new ConcurrentHashMap<>();


    /**
     * 新增socket
     *
     * @param session
     * @throws Exception
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String clientId = WebSocketInterceptor.getClientId(session.getUri());
        if (clientId != null) {
            Set<WebSocketSession> sessions = USERS.computeIfAbsent(clientId, k -> new HashSet<>());
            sessions.add(session);
            session.sendMessage(new TextMessage("成功建立socket连接"));
        }
        log.info("{}成功建立连接,当前用户数：{}", clientId, USERS.size());
    }

    /**
     * 接收socket信息
     *
     * @param webSocketSession
     * @param webSocketMessage
     * @throws Exception
     */
    @Override
    public void handleMessage(WebSocketSession webSocketSession, WebSocketMessage<?> webSocketMessage) throws Exception {
        String clientId = null;
        try {
            clientId = getClientId(webSocketSession);
            sendMsgToUser(getClientId(webSocketSession), new TextMessage(webSocketMessage.getPayload().toString()));
        } catch (Exception e) {
            log.warn("发送消息{}至用户{}失败", webSocketMessage, clientId);
        }

    }

    /**
     * 发送信息给指定用户
     *
     * @param clientId
     * @param message
     * @return
     */
    public boolean sendMsgToUser(String clientId, TextMessage message) {
        Set<WebSocketSession> sessions = USERS.get(clientId);
        if (sessions == null || sessions.size() <= 0) {
            log.warn("不存在用户:{},发送消息失败", clientId);
            return false;
        }
        boolean state = false;
        for (WebSocketSession session : sessions) {
            if (!session.isOpen()) {
                sessions.remove(session);
                continue;
            }
            try {
                log.info("发送消息{}至用户:{}", clientId, message);
                session.sendMessage(message);
                state = true;
            } catch (IOException e) {
                e.printStackTrace();
            }
        }

        return state;
    }

    /**
     * 广播信息
     *
     * @param message
     * @return
     */
    public boolean boardcast(TextMessage message) {
        boolean allSendSuccess = true;
        Set<String> clientIds = USERS.keySet();
        for (String clientId : clientIds) {
            boolean b = sendMsgToUser(clientId, message);
            if (b == false) {
                allSendSuccess = false;
            }
        }
        return allSendSuccess;
    }


    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        if (session.isOpen()) {
            session.close();
        }
        log.warn("连接出错:{}", session);
        removeSession(session);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        log.info("连接已关闭：{},status:{}", session, status);
        removeSession(session);
    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }

    /**
     * 获取用户标识
     *
     * @param session
     * @return
     */
    private String getClientId(WebSocketSession session) {
        try {
            return WebSocketInterceptor.getClientId(session.getUri());
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        }
    }

    private void removeSession(WebSocketSession session) {
        String clientId = getClientId(session);
        synchronized (clientId) {
            Set<WebSocketSession> sessions = USERS.get(clientId);
            if (sessions != null) {
                sessions.remove(session);
            }
        }
    }
}